import pandas as pd
from gofishing.marketData import _price_volume_columns
from gofishing.marketData import _pvFromSQLiteDB
from gofishing.base import Base, BaseBacktest, BaseRealtimeTQ, Log


class CommodityBase:

    commodities: dict[str, 'CommodityBase'] = {}

    def __init__(self, commodityID: str, contractCode: str,
                 frequency: str, length: int,
                 marginRate: float = 0.1,
                 mult: int = 10,
                 openCommission: float = 10.0,
                 closeCommission: float = 20.0) -> None:
        if commodityID in CommodityBase.commodities:
            raise Exception('commodityID must be unique')
        else:
            self.commodityID = commodityID
            CommodityBase.commodities[self.commodityID] = self

        self.contractCode = contractCode
        self.frequency = frequency
        self.length = length
        self.marginRate = marginRate            # 保证金比率
        self.mult = mult                        # 交易单位
        self.openCommission = openCommission    # 开仓手续费
        self.closeCommission = closeCommission  # 平仓手续费

        self.pv = pd.DataFrame()
        self.open = pd.Series(dtype='float32')
        self.high = pd.Series(dtype='float32')
        self.low = pd.Series(dtype='float32')
        self.close = pd.Series(dtype='float32')
        self.volume = pd.Series(dtype='float64')

    def updatePV(self) -> None:
        raise NotImplementedError

    def onCloseOfDay(self) -> None:
        raise NotImplementedError

    def update() -> None:
        commodities = CommodityBase.commodities
        for commodityID in commodities:
            commodities[commodityID].updatePV()
        if Base.isCloseOfDay():
            for commodityID in commodities:
                commodities[commodityID].onCloseOfDay()

    def isDominant(self) -> bool:
        raise NotImplementedError

    def switchToDominant(self) -> None:
        raise NotImplementedError

    def findCommodity(contractCode: str) -> 'CommodityBase':
        '''
            同一个 contractCode 对应多个 commodity 时，
            只返回第一个对应的 commodity
        '''
        commodities = CommodityBase.commodities
        for commodityID in commodities:
            if commodities[commodityID].contractCode == contractCode:
                return commodities[commodityID]
        raise Exception('no proper commodity found')


class CommodityBacktest(CommodityBase):

    def __init__(self, commodityID: str,
                 frequency: str, frequencyHigh: str,
                 length: int,
                 contractCode: str = '',
                 contractCodes: list[str] = [],
                 marginRate: float = 0.1,
                 mult: int = 10,
                 openCommission: float = 10,
                 closeCommission: float = 20) -> None:
        '''
            如果 contractCodes 不空，则忽略传入的 contractCode ，其值由程序自主决定
        '''
        super().__init__(commodityID, contractCode, frequency,
                         length, marginRate, mult,
                         openCommission, closeCommission)
        self.frequencyHigh = frequencyHigh
        self.contractCodes = contractCodes
        self._switchingToDominant = ''
        self._contractSwitchTime: list[pd.Timestamp] = []
        self._contractSwitchCode: list[str] = []

    def _loadData(self) -> None:
        self._pvComplete = _pvFromSQLiteDB(code=self.contractCode,
                                           frequency=self.frequency)
        self._pvCompleteHigh = _pvFromSQLiteDB(code=self.contractCode,
                                               frequency=self.frequencyHigh)

    def _loadPVDay(self) -> None:
        self._pvDay: dict[str, pd.DataFrame] = {}
        for contractCode in self.contractCodes:
            self._pvDay[contractCode] = _pvFromSQLiteDB(code=contractCode,
                                                        frequency='1d')

    def _pvClip(self,
                start: pd.Timestamp = pd.Timestamp(0),
                end: pd.Timestamp = pd.Timestamp(0),
                highFrequency=False) -> pd.DataFrame:
        pv = self._pvComplete.copy()
        if highFrequency:
            pv = self._pvCompleteHigh.copy()
        if start == pd.Timestamp(0):
            bars = self.length
            pv = pv[:end][-bars:]
            if len(pv) != bars:
                raise Exception('no enough data')
        else:
            pv = pv[start:end]
        return pv

    def _newBarTime(self) -> pd.Timestamp:
        now = Base.now()
        index = self._pvComplete.index
        for time in index:
            if now <= time:
                return time
        raise Exception('time out of index')

    def _dominantContract(self) -> str:
        dominantContract = ''
        pvDay = self._pvDay
        nowDate = pd.Timestamp(Base.now().date())
        for contractCode in pvDay:
            pv = pvDay[contractCode]
            if nowDate in pv.index:
                if not dominantContract:
                    dominantContract = contractCode
                elif pv.at[nowDate, 'volume'] > \
                        pvDay[dominantContract].at[nowDate, 'volume']:
                    dominantContract = contractCode
        return dominantContract

    def updatePV(self) -> None:
        if BaseBacktest.isFirstUpdate():
            if self.contractCodes:
                self._loadPVDay()
                dominantContract = self._dominantContract()
                if dominantContract:
                    self.contractCode = dominantContract
                else:
                    raise Exception('dominantContract error')
            self._loadData()

        if self._switchingToDominant:
            self.contractCode = self._switchingToDominant
            self._switchingToDominant = ''
            self._contractSwitchTime.append(BaseBacktest.now())
            self._contractSwitchCode.append(self.contractCode)
            self._loadData()

        newBarTime = self._newBarTime()
        self.pv = self._pvClip(end=newBarTime)

        # 把小周期多出大周期的那部分累加上去，作为大周期末尾时间戳的动态更新
        now = Base.now()
        if now != newBarTime:
            oldBarTime = self.pv.index[-2]
            highFrequencyPV = self._pvClip(
                start=oldBarTime, end=now, highFrequency=True)
            if highFrequencyPV.index[0] != oldBarTime:
                raise Exception('time index not aligned')
            highFrequencyPV = highFrequencyPV[1:]
            self.pv.at[newBarTime, 'open'] = highFrequencyPV['open'][0]
            self.pv.at[newBarTime, 'high'] = highFrequencyPV['high'].max()
            self.pv.at[newBarTime, 'low'] = highFrequencyPV['low'].min()
            self.pv.at[newBarTime, 'close'] = highFrequencyPV['close'][-1]
            self.pv.at[newBarTime, 'volume'] = highFrequencyPV['volume'].sum()

        self.open = self.pv['open']
        self.high = self.pv['high']
        self.low = self.pv['low']
        self.close = self.pv['close']
        self.volume = self.pv['volume']

    def onCloseOfDay(self) -> None:
        if self.contractCodes:
            if self._dominantContract() == self.contractCode:
                self._isDominant = True
            else:
                self._isDominant = False
        else:
            self._isDominant = True

    def isDominant(self) -> bool:
        return self._isDominant

    def switchToDominant(self) -> None:
        dominantContract = self._dominantContract()
        if self.contractCode == dominantContract:
            Log.warning('contract already dominant')
        else:
            self._switchingToDominant = dominantContract
            self._contractSwitchTime.append(BaseBacktest.now())
            self._contractSwitchCode.append(self.contractCode)

    def oneBarPV(self, start: pd.Timestamp, end: pd.Timestamp) -> pd.DataFrame:
        bars = self._pvCompleteHigh[start:end][1:]
        bar = bars[-1:].copy()
        bar.at[end, 'open'] = bars['open'][0]
        bar.at[end, 'high'] = bars['high'].max()
        bar.at[end, 'low'] = bars['low'].min()
        bar.at[end, 'close'] = bars['close'][-1]
        bar.at[end, 'volume'] = bars['volume'].sum()
        return bar

    def continuousPV(self, start: pd.Timestamp, end: pd.Timestamp,
                     highFrequency=False) -> pd.DataFrame:
        frequency = self.frequency
        if highFrequency:
            frequency = self.frequencyHigh
        contractSwitchTime = self._contractSwitchTime
        contractSwitchCode = self._contractSwitchCode
        if contractSwitchTime:
            selectedTime: list[pd.Timestamp] = []
            selectedCode: list[str] = []
            for time in contractSwitchTime:
                if start <= time and time <= end:
                    selectedTime.append(time)
                    selectedCode.append(
                        contractSwitchCode[contractSwitchTime.index(time)])
            if selectedCode[0] != selectedCode[1]:
                selectedTime.insert(0, start)
                selectedCode.insert(0, selectedCode[0])
            if selectedCode[-1] != selectedCode[-2]:
                selectedTime.append(end)
                selectedCode.append(selectedCode[-1])
            pvs: list[pd.DataFrame] = []
            indice = 0
            maxIndice = len(selectedTime) - 1
            while indice <= maxIndice:
                pvs.append(
                    _pvFromSQLiteDB(
                        code=selectedCode[indice],
                        frequency=frequency,
                        start=selectedTime[indice].isoformat(sep=' '),
                        end=selectedTime[indice+1].isoformat(sep=' ')))
                indice += 2
            return pd.concat(pvs)
        else:
            return _pvFromSQLiteDB(
                code=self.contractCode,
                frequency=frequency,
                start=start.isoformat(sep=' '),
                end=end.isoformat(sep=' '))


class CommodityRealtimeTQ(CommodityBase):

    def __init__(self, commodityID: str, contractCode: str,
                 frequency: str,
                 length: int,
                 marginRate: float = 0.1,
                 mult: int = 10,
                 openCommission: float = 10,
                 closeCommission: float = 20) -> None:
        super().__init__(commodityID, contractCode, frequency,
                         length, marginRate, mult,
                         openCommission, closeCommission)
        self.tqApi = BaseRealtimeTQ.tqApi
        self._kline = \
            self.tqApi.get_kline_serial(
                symbol=self.contractCode,
                duration_seconds=pd.Timedelta(self.frequency).seconds,
                data_length=self.length)
        self._timeCorrection = \
            pd.Timestamp('1970-01-01 08:00:00') + pd.Timedelta(self.frequency)
        pass

    def updatePV(self) -> None:
        kline = self._kline.copy()
        kline['time'] = \
            pd.to_datetime(kline['datetime'], origin=self._timeCorrection)
        kline.set_index('time', inplace=True)
        kline['code'] = kline['symbol']
        kline['frequency'] = self.frequency
        self.pv = kline[_price_volume_columns]

        self.open = self.pv['open']
        self.high = self.pv['high']
        self.low = self.pv['low']
        self.close = self.pv['close']
        self.volume = self.pv['volume']
